"""
# 公众号：悠悠智汇笔记
Author: 悠悠智汇笔记
Created: 2025-06-10
"""
# === 引入常用的可视化和数据处理库 ===
import pandas as pd  # 用于读取和处理表格数据
import numpy as np  # 用于数值处理，如处理缺失值
import seaborn as sns  # 用于绘制统计图表，这里用于绘制 clustermap 热图
import matplotlib.pyplot as plt  # matplotlib 的主绘图库
from matplotlib.patches import Patch  # 用于自定义图例中的色块
from matplotlib.colors import LinearSegmentedColormap, TwoSlopeNorm  # 用于自定义颜色映射
from mpl_toolkits.axes_grid1.inset_locator import inset_axes  # 用于插入色条到热图中

# === 1. 读取数据 ===
df = pd.read_csv("data.csv")  # 读取 CSV 文件，请替换为你自己的路径
df = df.set_index("#SampleID")  # 将第一列 SampleID 设置为行索引

# === 2. 提取标签和数值 ===
labels = df["Label"]  # 提取样本对应的标签列（用于分组配色）
data = df.drop("Label", axis=1)  # 删除标签列，保留纯数值数据
data = data.apply(pd.to_numeric, errors='coerce')  # 将所有值转换为数值型，非数字转为 NaN

# === 3. 缺失值填充 + 标准化 ===
data_filled = data.replace([np.inf, -np.inf], np.nan).fillna(data.mean())  # 将正负无穷替换为 NaN，并用平均值填充
standardized = (data_filled - data_filled.mean()) / data_filled.std()  # 标准化处理（Z-score）

# === 4. 为不同组定义对应颜色（用于顶部色条） ===
group_colors = {
    "Tumor": "#eb144c",  # 肿瘤组 - 红色
    "Co-housed+Tumor+Stress": "#4065d1",  # 混住+肿瘤+压力组 - 蓝色
    "Tumor+Stress": "#40af50"  # 肿瘤+压力组 - 绿色
}
col_colors = labels.map(group_colors)  # 将标签映射为颜色，形成一列颜色列表

# === 5. 自定义热图颜色渐变带：蓝 → 灰蓝 → 红 ===
custom_cmap = sns.diverging_palette(250, 10, s=75, l=40, as_cmap=True)  # 使用 seaborn 创建颜色带
norm = TwoSlopeNorm(vmin=-3, vcenter=0, vmax=3)  # 设置颜色映射范围和中心值（0表示中性）

# === 6. 图例固定为三组颜色，即使某组在当前数据中缺失也强制显示 ===
custom_legend = [
    Patch(facecolor="#eb144c", label="Tumor"),
    Patch(facecolor="#4065d1", label="Co+Tumor+Stress"),
    Patch(facecolor="#40af50", label="Tumor+Stress")
]

# === 7. 指定需要高亮（加粗变红）的代谢物名 ===
highlight_metabolites = {"Stearoylcarnitine", "Linoleylcarnitine", "Acetate"}

# === 8. 开始绘制 clustermap 热图 ===
sns.set(font_scale=0.5)  # 设置字体缩放比例

# 创建 clustermap 图，不显示默认的 colorbar
g = sns.clustermap(
    standardized.T,  # 注意：此处是转置后的数据，行为代谢物，列为样本
    cmap=custom_cmap,  # 使用自定义颜色带
    norm=norm,  # 使用定义的颜色归一化
    center=0,  # 中心值设置为0，便于红蓝对比
    col_colors=col_colors,  # 顶部添加样本标签颜色条
    linewidths=0.5,  # 设置单元格之间的分隔线宽度
    linecolor='#969687',  # 分隔线颜色
    xticklabels=False,  # 不显示列标签（样本名）
    yticklabels=True,  # 显示行标签（代谢物名）
    figsize=(9, 10),  # 图像大小（宽x高）
    dendrogram_ratio=(0.05, 0.03),  # 设置树状图相对大小（列树、行树）
    cbar_pos=None  # 禁用默认色条（后面手动添加）
)

# 设置树状图线条粗细（提高可读性）
for ax in [g.ax_row_dendrogram, g.ax_col_dendrogram]:
    for line in ax.collections:
        line.set_linewidth(1.2)

# 在热图右上角插入 colorbar 色条
cax = inset_axes(g.ax_heatmap,
                 width="3%",  # 色条宽度
                 height="25%",  # 色条高度
                 loc='upper left',  # 放置在热图的左上角
                 bbox_to_anchor=(1.3, 0.05, 0.5, 0.9),  # 相对位置微调
                 bbox_transform=g.ax_heatmap.transAxes,  # 使用热图坐标变换
                 borderpad=0)  # 边缘留白设置为0

# === 同步主图整体边距，防止元素挤压 ===
plt.subplots_adjust(right=0.08)  # 给右边留出空间放色条

# 手动添加色条
sm = plt.cm.ScalarMappable(cmap=custom_cmap, norm=norm)  # 创建色条映射器
sm.set_array([])  # 设置空数组用于绘制（不实际映射数据）
cbar = plt.colorbar(sm, cax=cax, orientation='vertical')  # 添加色条到指定 inset 坐标轴
cbar.ax.tick_params(labelsize=8)  # 设置色条刻度字体大小

# 设置颜色条标签字体大小（如果存在 "Label" 标签）
for ax in g.fig.axes:
    for label in ax.get_yticklabels():
        if label.get_text() == "Label":
            label.set_fontsize(12)
            label.set_fontweight("bold")

# === 9. 在右上角添加分组图例 ===
g.ax_col_dendrogram.legend(
    handles=custom_legend,  # 使用自定义图例色块
    loc="upper left",  # 图例放在右上角
    bbox_to_anchor=(1, 3),  # 图例框定位
    frameon=False,  # 不显示图例边框
    fontsize=10  # 图例字体大小
)

# === 10. 将需要高亮的代谢物标签加粗变红 ===
for label in g.ax_heatmap.get_yticklabels():  # 遍历所有 y 轴标签（代谢物名）
    label.set_fontsize(10)  # 设置字体大小
    if label.get_text() in highlight_metabolites:
        label.set_fontweight('bold')  # 加粗字体
        label.set_color('red')  # 设置为红色

# === 11. 最后整体调整布局 & 保存图像 ===
plt.subplots_adjust(right=0.75, top=0.92)  # 进一步微调图像边距，防止内容重叠
plt.savefig('聚类热图.png', dpi=300, bbox_inches='tight')  # 保存为高清图像文件
plt.show()  # 显示图像
